import torch
import torch.nn as nn

class CAM(nn.Module):
    # eca
    def __init__(self, channel, k_size=3):
        super(CAM, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # feature descriptor on the global spatial information
        y = self.avg_pool(x)

        # Two different branches of ECA module
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion
        y = self.sigmoid(y)
        x = x * y.expand_as(x)

        return x * y.expand_as(x)
    
class SAM(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.Conv1x1 = nn.Conv2d(in_channels, 1, kernel_size=1, bias=False)
        self.norm = nn.Sigmoid()

    def forward(self, U):
        q = self.Conv1x1(U)  # U:[bs,c,h,w] to q:[bs,1,h,w]
        q = self.norm(q)
        return U * q

class CBAM1(nn.Module):
    def __init__(self, c1,c2):
        super().__init__()
        self.cam = CAM(c1)
        self.sam = SAM(c1)

    def forward(self, U):
        U_sam = self.sam(U)
        U_cam = self.cam(U)
        return U_cam+U_sam